import parser
import os
from .base import Base

class PytorchDynamicEx2ONNX(Base):
    def __init__(self, success_list, failed_list):
        print('PytorchDynamicEx2ONNX Init--------------------------')
        self.tool_path = os.path.join(parser.get_value('maca_converter_root'), parser.get_value('maca_converter_path'))
        self.model_path = os.path.join(parser.get_value('model_root'), parser.get_value('pytorch2onnx_dynamic_model_path2'))
        self.model_class_name = parser.get_value('pytorch2onnx_dynamic_model_class_name2')
        self.input_shape = parser.get_value('pytorch2onnx_dynamic_model_input_shape2')
        self.model_weights_file = os.path.join(parser.get_value('model_root'), parser.get_value('pytorch2onnx_dynamic_model_weights_file2'))
        self.success_list = success_list
        self.failed_list = failed_list
        self.input_npy = os.path.join(parser.get_value('model_root'), parser.get_value('pytorch2onnx_dynamic_model_input_npy2'))
        self.output_npy = os.path.join(parser.get_value('model_root'), parser.get_value('pytorch2onnx_dynamic_model_output_npy2'))


    def test(self):
        cwd = os.getcwd()
        output = os.path.join(cwd, 'tmp.onnx')

        cmd = 'python ' + 'model_convert.py'  + ' --model_path ' + self.model_path + ' --model_type pytorch ' + '--output ' + output + \
                ' --model_class_name ' + self.model_class_name + \
                ' --input_shape ' + self.input_shape  + ' --model_weights_file ' + self.model_weights_file
                
        
        print('test_pytorch_dynamic2_2_onnx: ', cmd)

        self.test_model2onnx(self.tool_path, 'pytorch_dynamic2_2_onnx', cmd, output, self.success_list, self.failed_list, self.input_npy, self.output_npy)